package com.study.chapter09;

import com.study.chapter05.source.ClickSource;
import com.study.entity.Event;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

/**
 * @Description:
 * @Author: LiuQun
 * @Date: 2022/8/16 22:15
 */
public class BufferingSinkExample {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(
                        WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                                .withTimestampAssigner((event, l) -> event.timestamp)
                );
        stream.print("input==>");

        stream.addSink(new BufferingSink(10));

        env.execute();
    }

    public static class BufferingSink implements SinkFunction<Event>, CheckpointedFunction{
        //定义阈值
        private int threshold;
        //缓存数据集合
        private List<Event> bufferDataList;
        //定义一个算子状态
        private ListState<Event> checkpointState;

        public BufferingSink(int threshold) {
            this.threshold = threshold;
            bufferDataList = new ArrayList<>(threshold);
        }

        @Override
        public void invoke(Event value, Context context) throws Exception {
            //每来一条数据，则缓存到列表
            bufferDataList.add(value);
            //数据达到阈值，则入库
            if (bufferDataList.size() == threshold){
                //输出到控制台模拟入库
                bufferDataList.forEach(System.out::println);
                System.out.println("=============输出结束=================");
                //清空列表
                bufferDataList.clear();
            }
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            //进行快照处理
            checkpointState.clear(); //清空状态
            //复制缓存的列表数据到列表状态
            for (Event event : bufferDataList) {
                checkpointState.add(event);
            }
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            //初始化算子状态
            ListStateDescriptor<Event> descriptor = new ListStateDescriptor<>("buffer-state", Event.class);
            checkpointState = context.getOperatorStateStore().getListState(descriptor);

            //若从故障中恢复，则要将ListState中的所有元素复制到列表中
            if (context.isRestored()){
                checkpointState.get().forEach(event -> bufferDataList.add(event));
            }
        }
    }
}
